package cn.lishiyuan.algorithm.heap;

import java.util.Arrays;

/**
 * 堆
 */
public class Heap<T extends Comparable<T>> {

    private T[] data;

    private int size;

    private int capacity;

    private HeapType type;

    public Heap(T[] data){
        if(data == null || data.length == 0){
            throw new IllegalArgumentException("data is null or empty");
        }
        this.data = Arrays.copyOf(data, data.length);
        this.capacity = data.length;
        this.size = capacity;
        this.type = HeapType.MIN;

        // 从最后一个非叶子节点开始堆化
        for (int i = (capacity - 1) / 2; i >= 0; i--) {
            heapifyDown(this.data,size,i);
        }
    }

    public Heap(int capacity){
        this(HeapType.MIN,capacity);
    }

    public Heap(HeapType type,int capacity) {
        // 保证2的n次方
        this.type = type;
        this.capacity = capacity;
        this.data = (T[]) new Comparable<?>[capacity];
        this.size = 0;
    }


    // 大顶堆还是小顶堆
    public enum HeapType {
        BIG,
        MIN
    }

    public int size() {
        return size;
    }

    public boolean isEmpty() {
        return size == 0;
    }

    public T find(int index) {
        if(index < 0 || size <= index){
            return null;
        }
        return data[index];
    }

    public void add(T element) {
        if (element == null) {
            throw new IllegalArgumentException("Element cannot be null");
        }

        if(size == capacity){
            return;
        }
        int index= size;
        size++;
        data[index] = element;
        // 堆化
        heapifyUp();
    }

    public void update(T element) {
        if (element == null) {
            throw new IllegalArgumentException("Element cannot be null");
        }
        // 找到对应元素更新
        for (int i = 0; i < size; i++) {
            if(data[i].equals(element)){
                data[i] = element;
                break;
            }
        }

        // 堆化
        heapifyUp();
    }

    public T top() {
        if(isEmpty()){
            return null;
        }
        return data[0];
    }

    public T removeTop() {
        // 删除堆顶元素，进行堆化
        if(isEmpty()){
            return null;
        }
        T top = data[0];
        data[0] = data[size - 1];
        data[size - 1] = null;
        size--;
        heapifyDown(this.data,size,0);
        return top;
    }

    // 删除堆顶元素后堆化
    private void heapifyDown(T[] data,int size,int index) {
        // 从该index进行自上而下的堆化
        /**
         *        [a1]
         *       /  \
         *   [b2]   [c3]
         *   /\      /\
         *[d4][e5][e6][f7]
         *
         * ```
         * index(L) = index(E) * 2 + 1
         * index(R) = index(E) * 2 + 2
         * index(P) = (index(E) - 1) / 2
         * ```
         */
        // 大顶堆
        if(type == HeapType.BIG){
            // 与每一层的子节点比较大小，找出大的那个与其交换位置
            while (true){
                int maxIndex = index;
                int lIndex = index * 2 + 1;
                int rIndex = index * 2 + 2;

                if(lIndex < size && data[lIndex].compareTo(data[maxIndex]) > 0){
                    // 左节点比较大
                    maxIndex = lIndex;
                }

                if(rIndex < size && data[rIndex].compareTo(data[maxIndex]) > 0){
                    // 右节点比较大
                    maxIndex = rIndex;
                }

                if(maxIndex == index){
                    // 当前节点不比任何节点大了
                    return;
                }
                T temp = data[index];
                data[index] = data[maxIndex];
                data[maxIndex] = temp;
                index = maxIndex;
            }

        }else {
            // 与每一层的子节点比较大小，找出小的那个与其交换位置
            while (true){
                int minIndex = index;
                int lIndex = index * 2 + 1;
                int rIndex = index * 2 + 2;

                if(lIndex < size && data[lIndex].compareTo(data[minIndex]) < 0){
                    // 左节点比较小
                    minIndex = lIndex;
                }

                if(rIndex < size && data[rIndex].compareTo(data[minIndex]) < 0){
                    // 右节点比较小
                    minIndex = rIndex;
                }

                if(minIndex == index){
                    // 当前节点不比任何节点小了
                    return;
                }
                T temp = data[index];
                data[index] = data[minIndex];
                data[minIndex] = temp;
                index = minIndex;
            }
        }
    }

    // 添加元素后堆化
    private void heapifyUp(){
        int index = size - 1;
        // 从该index进行自底下而上的堆化
        /**
         *        [a1]
         *       /  \
         *   [b2]   [c3]
         *   /\      /\
         *[d4][e5][e6][f7]
         *
         * ```
         * index(L) = index(E) * 2 + 1
         * index(R) = index(E) * 2 + 2
         * index(P) = (index(E) - 1) / 2
         * ```
         *
         */
        // 大顶堆
        if(type == HeapType.BIG){
            // index/2 父节点,当前节点比父节点大则交换位置
            int pIndex = (index-1) / 2;
            // -1 / 2 = 0
            while ( index !=0 && pIndex >= 0 && data[index].compareTo(data[pIndex]) > 0){
                T temp = data[index];
                data[index] = data[pIndex];
                data[pIndex] = temp;
                index = pIndex;
                pIndex = (index - 1) / 2;
            }
        }else {
            // index/2 父节点,当前节点比父节点小则交换位置
            int pIndex = (index-1) / 2;
            // -1 / 2 = 0
            while (index !=0 && pIndex >= 0 && data[index].compareTo(data[pIndex]) < 0){
                T temp = data[index];
                data[index] = data[pIndex];
                data[pIndex] = temp;
                index = pIndex;
                pIndex = (index-1) / 2;
            }
        }
    }

    public void clear() {
        size = 0;
        Arrays.fill(data, null);
    }


    public HeapType getType() {
        return type;
    }

    public boolean isFull() {
        return capacity == size;
    }
}
